from io import BytesIO

import torch
from PIL import Image
from flask import Flask, request
from rapid_table import RapidTable, RapidTableInput
from rapidocr import RapidOCR
from ultralytics import YOLO

from utils.detector import detect_objects

app = Flask(__name__)

app.config.from_pyfile('config.py')
use_cuda = torch.cuda.is_available()
device = 'cuda:0' if use_cuda else 'cpu'

model_path = 'yolov11m_phase_epoches100.pt'
model = YOLO(model_path)
# 模型量化（可选）
if use_cuda:
    model = model.to('cuda:0')
    model = model.half()  # 转为半精度浮点数
else:
    model = model.to('cpu')

ocr_engine = RapidOCR()
input_args = RapidTableInput(model_type="unitable", use_cuda=use_cuda, device=device)
table_engine = RapidTable(input_args)


@app.route('/')
def hello_world():
    return 'Hello World!'


@app.route('/predict', methods=['POST'])
def predict():
    if 'file' not in request.files:
        return '未包含文件部分'
    image = request.files['file']
    if image.filename == '':
        return '没有选择文件'
    if image:
        try:
            img = Image.open(BytesIO(image.read()))
            # 使用 with 语句确保上下文清理
            with torch.no_grad():
                result = detect_objects(img, model, ocr_engine, table_engine)
            return result
        except Exception as e:
            print(e)
            return 'error'
    else:
        return "未收到文件"


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)
